//! JIT与VM的集成模块
//! 
//! 提供JIT编译器与Pipit虚拟机的无缝集成

use std::collections::HashMap;
use std::time::{Duration, Instant};


use crate::jit::{PipitJIT, JITOptimizationLevel, CompilationError};

use crate::vm::VirtualMachine;

/// JIT执行上下文
pub struct JITContext {
    compiler: PipitJIT,
    compiled_functions: HashMap<String, *const u8>,
    execution_stats: ExecutionStats,
}

/// 执行统计
#[derive(Debug, Default)]
pub struct ExecutionStats {
    pub total_compilations: u64,
    pub total_executions: u64,
    pub total_compilation_time: Duration,
    pub total_execution_time: Duration,
    pub cache_hits: u64,
    pub cache_misses: u64,
}

/// JIT配置
#[derive(Debug, Clone)]
pub struct JITConfig {
    pub enabled: bool,
    pub optimization_level: JITOptimizationLevel,
    pub max_cache_size: usize,
    pub enable_profiling: bool,
    pub inline_threshold: u32,
}

impl Default for JITConfig {
    fn default() -> Self {
        Self {
            enabled: true,
            optimization_level: JITOptimizationLevel::Aggressive,
            max_cache_size: 64 * 1024 * 1024, // 64MB
            enable_profiling: false,
            inline_threshold: 200,
        }
    }
}

impl JITContext {
    /// 创建新的JIT上下文
    pub fn new(enabled: bool, optimization_level: JITOptimizationLevel) -> Self {
        let config = JITConfig {
            enabled,
            optimization_level,
            ..Default::default()
        };
        
        Self::with_config(config)
    }

    /// 使用配置创建JIT上下文
    pub fn with_config(config: JITConfig) -> Self {
        Self {
            compiler: PipitJIT::new(config.optimization_level),
            compiled_functions: HashMap::new(),
            execution_stats: ExecutionStats::default(),
        }
    }

    /// 编译函数到机器码
    pub fn compile_function(&mut self, _vm: &mut VirtualMachine, name: &str, bytecode: &[u8]) -> Result<(), CompilationError> {
        let start = Instant::now();
        
        let func_ptr = self.compiler.compile_function(name, bytecode)?;
        self.compiled_functions.insert(name.to_string(), func_ptr);
        
        self.execution_stats.total_compilations += 1;
        self.execution_stats.total_compilation_time += start.elapsed();
        
        Ok(())
    }

    /// 执行已编译的函数
    pub fn execute_function(&mut self, _vm: &mut VirtualMachine, name: &str, args: &[i64]) -> Option<i64> {
        let start = Instant::now();
        
        let func_ptr = *self.compiled_functions.get(name)?;
        
        // 执行机器码
        let result = unsafe {
            let func: extern "C" fn(i64, i64, i64, i64) -> i64 = std::mem::transmute(func_ptr);
            func(args[0], args[1], args[2], args[3])
        };
        
        self.execution_stats.total_executions += 1;
        self.execution_stats.total_execution_time += start.elapsed();
        
        Some(result)
    }

    /// 获取执行统计
    pub fn get_stats(&self) -> &ExecutionStats {
        &self.execution_stats
    }

    /// 清除缓存
    pub fn clear_cache(&mut self) {
        self.compiled_functions.clear();
        self.execution_stats = ExecutionStats::default();
    }

    /// 获取缓存大小
    pub fn cache_size(&self) -> usize {
        self.compiled_functions.len()
    }

    /// 设置优化级别
    pub fn set_optimization_level(&mut self, level: JITOptimizationLevel) {
        // 创建新的JIT实例以应用新的优化级别
        self.compiler = PipitJIT::new(level);
    }
}

/// JIT执行环境
pub struct JITExecutionContext {
    jit_ctx: JITContext,
    vm: VirtualMachine,
}

impl JITExecutionContext {
    pub fn new(config: JITConfig) -> Self {
        Self {
            jit_ctx: JITContext::with_config(config),
            vm: VirtualMachine::new(),
        }
    }

    /// 运行脚本
    pub fn run_script(&mut self, _script: &str) -> Result<i64, String> {
        // 在实际实现中，这里会解析和编译脚本
        // 简化版本：返回模拟结果
        Ok(42)
    }

    /// 运行字节码
    pub fn run_bytecode(&mut self, name: &str, bytecode: &[u8], args: &[i64]) -> Result<i64, String> {
        self.jit_ctx.compile_function(&mut self.vm, name, bytecode);
        
        match self.jit_ctx.execute_function(&mut self.vm, name, args) {
            Some(result) => Ok(result),
            None => Err(format!("Function '{}' not found", name)),
        }
    }
}

/// JIT性能分析器
pub struct JITProfiler {
    measurements: HashMap<String, Vec<Duration>>,
}

impl JITProfiler {
    pub fn new() -> Self {
        Self {
            measurements: HashMap::new(),
        }
    }

    /// 开始测量
    pub fn start_measurement(&mut self, name: &str) -> MeasurementHandle {
        MeasurementHandle {
            name: name.to_string(),
            start: Instant::now(),
            profiler: self,
        }
    }

    /// 记录测量结果
    fn record_measurement(&mut self, name: &str, duration: Duration) {
        self.measurements
            .entry(name.to_string())
            .or_insert_with(Vec::new)
            .push(duration);
    }

    /// 获取平均执行时间
    pub fn get_average_time(&self, name: &str) -> Option<Duration> {
        self.measurements
            .get(name)
            .and_then(|measurements| {
                if measurements.is_empty() {
                    None
                } else {
                    let sum: Duration = measurements.iter().sum();
                    Some(sum / measurements.len() as u32)
                }
            })
    }

    /// 清除所有测量
    pub fn clear(&mut self) {
        self.measurements.clear();
    }
}

/// 测量句柄
pub struct MeasurementHandle<'a> {
    name: String,
    start: Instant,
    profiler: &'a mut JITProfiler,
}

impl<'a> Drop for MeasurementHandle<'a> {
    fn drop(&mut self) {
        let duration = self.start.elapsed();
        self.profiler.record_measurement(&self.name, duration);
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_jit_context_creation() {
        let jit_ctx = JITContext::new(true, JITOptimizationLevel::Basic);
        assert!(jit_ctx.cache_size() == 0);
    }

    #[test]
    fn test_function_compilation() {
        let mut jit_ctx = JITContext::new(true, JITOptimizationLevel::Basic);
        let mut vm = VirtualMachine::new();
        
        let bytecode = vec![0x01, 0x02, 0x03];
        assert!(jit_ctx.compile_function(&mut vm, "test", &bytecode));
        assert!(jit_ctx.compiled_functions.contains_key("test"));
    }

    #[test]
    fn test_function_execution() {
        let mut jit_ctx = JITContext::new(true, JITOptimizationLevel::Basic);
        let mut vm = VirtualMachine::new();
        
        let bytecode = vec![0x01, 0x02, 0x03];
        jit_ctx.compile_function(&mut vm, "test", &bytecode);
        
        let result = jit_ctx.execute_function(&mut vm, "test", &[1, 2, 3]);
        assert!(result.is_some());
    }

    #[test]
    fn test_stats_collection() {
        let mut jit_ctx = JITContext::new(true, JITOptimizationLevel::Basic);
        let mut vm = VirtualMachine::new();
        
        let bytecode = vec![0x01, 0x02, 0x03];
        jit_ctx.compile_function(&mut vm, "test", &bytecode);
        
        let stats = jit_ctx.get_stats();
        assert!(stats.total_compilations > 0);
    }

    #[test]
    fn test_profiler() {
        let mut profiler = JITProfiler::new();
        
        {
            let _handle = profiler.start_measurement("test");
            std::thread::sleep(Duration::from_millis(1));
        }
        
        let avg_time = profiler.get_average_time("test");
        assert!(avg_time.is_some());
        assert!(avg_time.unwrap() >= Duration::from_millis(1));
    }
}